import os
from pathlib import Path
import threading
from queue import Queue, Empty
import logging
from tqdm import tqdm
from PIL import Image
import traceback
import time
from transformers.image_utils import load_image
from transformers import  AutoProcessor
# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('image_check_error.log'),
        logging.StreamHandler()
    ]
)

# 配置
NUM_THREADS = 20
MIN_FILE_SIZE = 1024  # 1KB，小于这个大小认为是空文件
ckpt = "google/siglip2-so400m-patch14-384"
processor = AutoProcessor.from_pretrained(ckpt)
class ImageChecker:
    def __init__(self, image_dirs, num_threads=NUM_THREADS):
        self.image_dirs = [Path(d) for d in image_dirs]
        self.num_threads = num_threads
        self.queue = Queue()
        self.invalid_images = Queue()
        self.pbar = None
        self.total_files = 0
        self.processed_files = 0
        self.lock = threading.Lock()

    def is_valid_image(self, image_path):
        """检查图片是否有效"""
        try:
            # 检查文件大小
            if os.path.getsize(image_path) < MIN_FILE_SIZE:
                return False, "File too small"

            # 尝试打开图片
            with Image.open(image_path) as img:
                # 检查图片大小
                if img.size[0] < 10 or img.size[1] < 10:
                    return False, "Image dimensions too small"
                
                # 验证图片内容
                img.verify()
                iii = load_image(image_path._str)
                _ = processor(images=[iii], return_tensors="pt")
                # 尝试加载图片数据
                # img = img.load()
                
                # 检查是否是完全透明或完全黑色的图片
                if img.mode in ('RGBA', 'LA'):
                    # 转换为RGBA模式以检查透明度
                    rgba_img = img.convert('RGBA')
                    # 获取透明度通道
                    alpha = rgba_img.split()[-1]
                    if all(p == 0 for p in alpha.getdata()):
                        return False, "Completely transparent image"

                return True, "OK"

        except (IOError, SyntaxError, ValueError, OSError) as e:
            return False, str(e)
        except Exception as e:
            return False, f"Unexpected error: {str(e)}"

    def load_images(self):
        """加载所有图片路径到队列"""
        total = 0
        for dir_path in self.image_dirs:
            if not dir_path.exists():
                logging.warning(f"Directory not found: {dir_path}")
                continue
                
            for img_path in dir_path.glob('*.jpg'):
                self.queue.put(img_path)
                total += 1
                
        self.total_files = total
        self.pbar = tqdm(total=total, desc="Checking images")
        logging.info(f"Found {total} images to check")

    def check_worker(self):
        """工作线程函数"""
        while True:
            try:
                # 获取下一个任务
                img_path = self.queue.get(block=True, timeout=5)
                
                # 检查图片
                is_valid, reason = self.is_valid_image(img_path)
                
                if not is_valid:
                    self.invalid_images.put((img_path, reason))
                    try:
                        # 删除无效图片
                        os.remove(img_path)
                        logging.info(f"Deleted invalid image: {img_path} (Reason: {reason})")
                    except Exception as e:
                        logging.error(f"Failed to delete {img_path}: {str(e)}")

                # 更新进度
                with self.lock:
                    self.processed_files += 1
                    self.pbar.update(1)
                
                self.queue.task_done()
                
            except Empty:
                break  # 队列为空，退出线程
            except Exception as e:
                logging.error(f"Worker error: {str(e)}\n{traceback.format_exc()}")
                continue

    def start(self):
        """开始检查过程"""
        try:
            logging.info("Starting image check process...")
            
            # 加载图片路径
            self.load_images()
            
            if self.total_files == 0:
                logging.warning("No images found to check!")
                return
            
            # 创建线程池
            threads = []
            for _ in range(self.num_threads):
                thread = threading.Thread(target=self.check_worker)
                thread.start()
                threads.append(thread)
            
            # 等待所有线程完成
            for thread in threads:
                thread.join()
            
            # 保存无效图片列表
            invalid_count = self.invalid_images.qsize()
            if invalid_count > 0:
                logging.warning(f"Found {invalid_count} invalid images")
                with open('invalid_images.txt', 'w', encoding='utf-8') as f:
                    f.write("Image Path,Reason\n")
                    while not self.invalid_images.empty():
                        img_path, reason = self.invalid_images.get()
                        f.write(f"{img_path},{reason}\n")
            
            logging.info(f"Check completed! Processed {self.processed_files} images, found {invalid_count} invalid images")
            
        except Exception as e:
            logging.error(f"Error in check process: {str(e)}\n{traceback.format_exc()}")
        finally:
            if self.pbar:
                self.pbar.close()

def main():
    # 配置图片目录
    image_dirs = [
        # 'data/eval_images_q',
        '/home/tfj/datasets/image_retri10k/eval_images_o2'
    ]
    
    # 创建检查器实例并开始检查
    checker = ImageChecker(image_dirs, NUM_THREADS)
    checker.start()

if __name__ == "__main__":
    main() 